import time

from aimodel import tongyi_model
from src.tools.dms_operator import dmsoperator


def head_contents_extract(file_json):
    '''提取前200字文本'''
    contexts_context = ""
    for par in file_json["paragraphInfos"]:
        if par["text"]:
            contexts_context = contexts_context + par["text"] + "\n"
        if len(contexts_context)>200:
            return contexts_context
    for table in file_json["tableInfos"]:
        for table_cell in table["tableCellInfos"]:
            for par in table_cell["paragraphInfos"]:
                contexts_context = contexts_context + par["text"] + "\n"
            if len(contexts_context) > 200:
                return contexts_context
    return contexts_context

def detect(contexts):
    if not contexts:
        return "其他"
    detect_prompt = f"已知以下文本是一段文本中的封面页；判断该段文本属于(生产批记录、工艺规程、工艺信息表、供应商目录、其他)中的哪一类；一般来说封面包含“记录”两字的为生产批记录；只用回答类别。\n===\n{contexts}\n===\n"
    lb = tongyi_model.invoke(detect_prompt)
    return lb.content

def detect_api(file_json):
    '''传入解析json；返回文件类别；枚举类型：（生产批记录、工艺规程、工艺信息表、供应商目录、其他）'''
    start_time = time.time()
    contexts = head_contents_extract(file_json)
    detect_result = detect(contexts)
    print(detect_result)
    end_time = time.time()
    print("运行时间: {:.6f} 秒".format(end_time - start_time))
    return detect_result



if __name__ == "__main__":
    file_json = dmsoperator.get_jsonfile("Files\\202412\\20241208103351Gif8BrplF68e636bc01f43f2a257e15065683a1b05.docx")
    start_time = time.time()
    contexts = head_contents_extract(file_json)
    print(detect(contexts))
    end_time = time.time()
    print("运行时间: {:.6f} 秒".format(end_time - start_time))